import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import yfinance as yf

# ==================
# 数据预处理模块
# ==================
def load_data(symbol="BTC-USD"):
    # 内置2024-2025年OHLC模拟数据生成器
    # dates = pd.date_range("2024-01-01", "2025-02-23")
    # return pd.DataFrame(
    #     {
    #         "Close": np.sin(np.linspace(0, 20, len(dates))) * 100 + 500,
    #         "High": np.sin(np.linspace(0, 20, len(dates))) * 120 + 520,
    #         "Low": np.sin(np.linspace(0, 20, len(dates))) * 80 + 480,
    #     },
    #     index=dates,
    # )

    # 获取股票信息
    data = yf.download(symbol,period="3wk",interval="1h")
    ohlc_data = data[["Open", "High", "Low", "Close"]]
    return ohlc_data


# ==================
# 指标计算核心
# ==================
class Strategy:
    @staticmethod
    def macd(close, fast=12, slow=26, signal=9):
        ema_fast = close.ewm(span=fast).mean()
        ema_slow = close.ewm(span=slow).mean()
        macd_line = ema_fast - ema_slow
        signal_line = macd_line.ewm(span=signal).mean()
        return macd_line, signal_line

    @staticmethod
    def kdj(high, low, close, n=9):
        min_low = low.rolling(n).min()
        max_high = high.rolling(n).max()
        rsv = (close - min_low) / (max_high - min_low) * 100
        K = rsv.ewm(com=2).mean()
        D = K.ewm(com=2).mean()
        J = 3 * K - 2 * D
        return K, D, J


# ==================
# 信号引擎
# ==================
def generate_signals(df):
    df["MACD"], df["Signal"] = Strategy.macd(df["Close"])
    df["K"], df["D"], df["J"] = Strategy.kdj(df["High"], df["Low"], df["Close"])

    # 多空信号逻辑
    df["Buy"] = (df["MACD"] > df["Signal"]) & (df["J"] < 30)
    df["Sell"] = (df["MACD"] < df["Signal"]) & (df["J"] > 70)
    return df


# ==================
# 专业可视化
# ==================
def plot_results(df):
    plt.figure(figsize=(16, 10))

    # 价格与信号层
    ax1 = plt.subplot(3, 1, 1)
    plt.plot(df["Close"], label="Price", color="navy")
    plt.scatter(
        df[df["Buy"]].index,
        df[df["Buy"]]["Close"],
        marker="^",
        color="red",
        s=100,
        label="Buy",
    )
    plt.scatter(
        df[df["Sell"]].index,
        df[df["Sell"]]["Close"],
        marker="v",
        color="green",
        s=100,
        label="Sell",
    )

    # MACD层
    plt.subplot(3, 1, 2, sharex=ax1)
    plt.plot(df["MACD"], label="MACD", color="darkorange")
    plt.plot(df["Signal"], label="Signal", color="forestgreen")
    plt.bar(
        df.index,
        df["MACD"] - df["Signal"],
        color=np.where(df["MACD"] > df["Signal"], "lime", "crimson"),
    )

    # KDJ层
    plt.subplot(3, 1, 3, sharex=ax1)
    plt.plot(df["K"], label="K", color="purple")
    plt.plot(df["D"], label="D", color="teal")
    plt.plot(df["J"], label="J", color="firebrick", linestyle="--")
    plt.axhline(80, color="gray", ls=":")
    plt.axhline(20, color="gray", ls=":")

    plt.tight_layout()
    plt.show()


# ==================
# 执行入口
# ==================
if __name__ == "__main__":
    df = load_data()
    df = generate_signals(df)
    plot_results(df)
